{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instructions for training the model for post-fault operation of the system.\n",
    "1. Uncomment Line 90 in openDSSenv34.py; which is L_OUT='L24'.\n",
    "2. Change the L_OUT to corresponding line failure from the candidate set to get results for those failures. Candidate_Lines=['L7','L9','L15','L16','L18','L19','L21','L22','L23','L24']\n",
    "3. Run the corresponding code blocks below to obtain results for PPO, A2C, and TRPO.\n",
    "4. Use CustomNN class to use Multilayer Perceptrons (MLP) as feature extractor, else use CustomCNN to use Convolution Neural Network (CNN as feature extractor)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gym\n",
    "from stable_baselines3 import PPO\n",
    "from stable_baselines3 import A2C\n",
    "from sb3_contrib import TRPO\n",
    "from openDSSenv34 import openDSSenv34\n",
    "import torch as th\n",
    "from stable_baselines3.common.utils import set_random_seed\n",
    "from numba import jit\n",
    "from state_action_reward import take_action, get_state\n",
    "\n",
    "\n",
    "from stable_baselines3.common.torch_layers import (\n",
    "    BaseFeaturesExtractor,\n",
    "    CombinedExtractor,\n",
    "    FlattenExtractor,\n",
    "    MlpExtractor,\n",
    "    NatureCNN,\n",
    "    create_mlp,\n",
    ")\n",
    "\n",
    "\n",
    "import pickle\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv\n",
    "\n",
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "import time\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare the Neural network for feature extraction. (1) MLP; (2) CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customer NN used in the report\n",
    "class CustomNN(BaseFeaturesExtractor): # MLP\n",
    "\n",
    "    \"\"\"\n",
    "    :param observation_space: (gym.Space)\n",
    "    :param features_dim: (int) Number of features extracted.\n",
    "        This corresponds to the number of unit for the last layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):\n",
    "        from torch import nn\n",
    "        super(CustomNN, self).__init__(observation_space, features_dim)\n",
    "\n",
    "\n",
    "        n_flatten = 1521# Need to adjust this to fit in the microgrid state space size\n",
    "        self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU())\n",
    "\n",
    "    def forward(self, observations):\n",
    "        \n",
    "        if len(observations[\"Unserved Energy\"].shape) == 1:\n",
    "            data_UE = observations[\"Unserved Energy\"][:, None]\n",
    "        else:\n",
    "            data_UE = observations[\"Unserved Energy\"]\n",
    "       \n",
    "        ################### Now Trying with the Unserved Energy into Account ################\n",
    "        statevec = np.concatenate((data_UE,\n",
    "                                    observations['NodeFeat(BusVoltage)'].flatten(1,2),\n",
    "                                    observations['EdgeFeat(branchflow)'][:,:],\n",
    "                                    observations['Adjacency'].flatten(1,2)), axis=1)\n",
    "        \n",
    "        statevec = np.array(statevec)\n",
    "        statevec = th.from_numpy(statevec)\n",
    "\n",
    "        # print(statevec)\n",
    "        return self.linear(statevec)\n",
    "    \n",
    "    \n",
    "class CustomCNN(BaseFeaturesExtractor): #CNN\n",
    "\n",
    "    \"\"\"\n",
    "    :param observation_space: (gym.Space)\n",
    "    :param features_dim: (int) Number of features extracted.\n",
    "        This corresponds to the number of unit for the last layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):\n",
    "        from torch import nn\n",
    "        super(CustomNN, self).__init__(observation_space, features_dim)\n",
    "\n",
    "\n",
    "        n_flatten = 1521# Need to adjust this to fit in the microgrid state space size\n",
    "\n",
    "        # define cnn layer for feature extraction\n",
    "        self.cnn_layers = nn.Sequential(\n",
    "            nn.Conv1d(1, 100, kernel_size=5, stride=3, padding=1),  # 1st 1D-CNN layer\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(100, 100, kernel_size=5, stride=3, padding=1),  # 2nd 1D-CNN layer\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(100, 100, kernel_size=5, stride=3, padding=1),  # 3rd 1D-CNN layer\n",
    "            nn.ReLU(),\n",
    "            nn.Flatten(1,-1),\n",
    "        )\n",
    "        # calculate the output shape from cnn_layers\n",
    "        with th.no_grad():\n",
    "            n_flatten = self.cnn_layers(\n",
    "                th.as_tensor(np.empty((1,1,n_flatten))).float()\n",
    "            ).shape[1]\n",
    "\n",
    "        # add a linear layer to get expected feature dimention\n",
    "        self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU())\n",
    "\n",
    "\n",
    "    def forward(self, observations):\n",
    "        # get selected observations as state vector\n",
    "        if len(observations[\"Unserved Energy\"].shape) == 1:\n",
    "            data_UE = observations[\"Unserved Energy\"][:, None]\n",
    "        else:\n",
    "            data_UE = observations[\"Unserved Energy\"]\n",
    "\n",
    "\n",
    "        ################### Now Trying with the Unserved Energy into Account ################\n",
    "        statevec = np.concatenate((data_UE,\n",
    "                                    observations['NodeFeat(BusVoltage)'].flatten(1,2),\n",
    "                                    observations['EdgeFeat(branchflow)'][:,:],\n",
    "                                    observations['Adjacency'].flatten(1,2)), axis=1)\n",
    "\n",
    "        statevec = np.array(statevec)\n",
    "        statevec = th.from_numpy(statevec)\n",
    "\n",
    "        # Add additional dimention to match 1d-cnn layer input shape.\n",
    "        statevec = statevec.unsqueeze(0)\n",
    "        statevec = statevec.transpose(1,0)\n",
    "\n",
    "    \n",
    "        return self.linear(self.cnn_layers(statevec))\n",
    "\n",
    "def make_env(rank, seed=0):\n",
    "    \"\"\"\n",
    "    Utility function for multiprocessed env.\n",
    "\n",
    "    :param env_id: (str) the environment ID\n",
    "    :param num_env: (int) the number of environments you wish to have in subprocesses\n",
    "    :param seed: (int) the inital seed for RNG\n",
    "    :param rank: (int) index of the subprocess\n",
    "    \"\"\"\n",
    "    def _init():\n",
    "        env = openDSSenv34()\n",
    "        env.seed(seed + rank)\n",
    "        return env\n",
    "    set_random_seed(seed)\n",
    "    return _init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(a) Run the following code for RL algorithm: PPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the model for PPO\n",
    "num_cpu = 8\n",
    "env=make_vec_env(openDSSenv34,n_envs=num_cpu,seed=0)\n",
    "\n",
    "rms_prop_eps = 1e-5\n",
    "\n",
    "# Training the model for PPO\n",
    "policy_kwargs = dict(\n",
    "    features_extractor_class=CustomNN, # chage with CustomCNN to use CNN as feature extractor\n",
    "    features_extractor_kwargs=dict(features_dim=256),\n",
    "    activation_fn=th.nn.Tanh,\n",
    "    net_arch=[dict(vf=[128, 128])] # \n",
    ")\n",
    "\n",
    "\n",
    "model = PPO('MultiInputPolicy', env,tensorboard_log=\"logger_PPO_narrow/\", policy_kwargs=policy_kwargs, verbose=1, n_steps=100, batch_size=100,\n",
    "            gamma=1.00,\n",
    "            learning_rate=0.000001,#0.00001 from 0.003 to 5e-6♠\n",
    "                ent_coef=0.01# 0.05\n",
    "                ).learn(total_timesteps=50000, n_eval_episodes=1, log_interval=1, tb_log_name=\"R1_Microgrid_env_mlp_Normal\")\n",
    "\n",
    "log_dir = \".\"\n",
    "model.save(log_dir + \"r1_MG_bus_mlp_with_entropy_05_multi_env_normal_PPO\"+str(datetime.datetime.now().day)+\"_\"+str(datetime.datetime.now().hour)+\"_\"+str(datetime.datetime.now().minute))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(b) Run the following code for RL algorithm: A2C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cpu = 8\n",
    "env=make_vec_env(openDSSenv34,n_envs=num_cpu,seed=0)\n",
    "\n",
    "rms_prop_eps = 1e-5\n",
    "\n",
    "# Training the model for A2C\n",
    "policy_kwargs = dict(\n",
    "    features_extractor_class=CustomNN, # chage with CustomCNN to use CNN as feature extractor\n",
    "    features_extractor_kwargs=dict(features_dim=128),\n",
    "    optimizer_class = th.optim.RMSprop,\n",
    "    optimizer_kwargs = dict(alpha=0.89, eps=rms_prop_eps, weight_decay=0)\n",
    ")\n",
    "model = A2C('MultiInputPolicy', env,tensorboard_log=\"logger/\", policy_kwargs=policy_kwargs, verbose=1, n_steps=100,\n",
    "        use_rms_prop=False,\n",
    "            gamma=1.00,\n",
    "            learning_rate=0.000001,\n",
    "            ).learn(total_timesteps=20000, n_eval_episodes=1, log_interval=1, tb_log_name=\"R1_A2C\")\n",
    "\n",
    "log_dir = \".\"\n",
    "model.save(log_dir + \"r1_MG_bus_mlp_with_entropy_05_multi_env_normal_A2C\"+str(datetime.datetime.now().day)+\"_\"+str(datetime.datetime.now().hour)+\"_\"+str(datetime.datetime.now().minute))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(c) Run the following code for RL algorithm: TRPO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cpu = 8\n",
    "env=make_vec_env(openDSSenv34,n_envs=num_cpu,seed=0)\n",
    "\n",
    "rms_prop_eps = 1e-5\n",
    "\n",
    "# Training the model for TRPO\n",
    "policy_kwargs = dict(\n",
    "    features_extractor_class=CustomNN, # chage with CustomCNN to use CNN as feature extractor\n",
    "    features_extractor_kwargs=dict(features_dim=128),\n",
    "    optimizer_class = th.optim.RMSprop,\n",
    "    optimizer_kwargs = dict(alpha=0.89, eps=rms_prop_eps, weight_decay=0)\n",
    ")\n",
    "\n",
    "model = TRPO('MultiInputPolicy', env,tensorboard_log=\"logger/\", policy_kwargs=policy_kwargs, verbose=1, n_steps=200, \n",
    "            gamma=1.00,\n",
    "            learning_rate=0.000001,\n",
    "            batch_size=100).learn(total_timesteps=80000, log_interval=1, tb_log_name=\"R1_Microgrid-TRPO\")\n",
    "\n",
    "log_dir = \".\"\n",
    "model.save(log_dir + \"r1_MG_TRPO\"+str(datetime.datetime.now().day)+\"_\"+str(datetime.datetime.now().hour)+\"_\"+str(datetime.datetime.now().minute))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training is done. Now, We can test it for post-disaster condition with following changes.\n",
    "1. Change line 73 in DSS_Initialize.py for testing each scenario, which is 'factor = 0.5'.\n",
    "2. Set the 'factor' to 0.5, 1, 1.5 and rerun following code.\n",
    "3. Collect the result to generate Table II in the report."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing Microgrid env with sectionalizing and tie switches\n",
      "Env initialized\n",
      "The loss is: 0.021264677755010977\n",
      "The topology violation status is: 200\n",
      "The voltage violation status is:  0\n",
      "The branch flow violation status is:  0\n",
      "The amount of unserved energy is: 0.5676008991317862\n",
      "The convergence status is: 0\n",
      "For Unity Load: The Optimal Configuration is :.{} tensor([[1., 0., 0., 1., 1., 0., 1., 0., 1.]])\n",
      "Run time [s]:  0.03441190719604492\n"
     ]
    }
   ],
   "source": [
    "import openDSSenv34\n",
    "import importlib\n",
    "importlib.reload(openDSSenv34)\n",
    "from openDSSenv34 import openDSSenv34\n",
    "\n",
    "env = openDSSenv34()\n",
    "obs, DSSCKTOBJ, G_INIT = env.new_test_func()\n",
    "start = time.time()\n",
    "obs = {key: th.as_tensor([_obs]) for (key, _obs) in obs.items()}\n",
    "obs['loss'] = th.as_tensor([[obs['loss']]])\n",
    "#obs['TopologicalConstr'] = torch.as_tensor([[obs['TopologicalConstr']]])\n",
    "obs['VoltageViolation'] = th.as_tensor([[obs['VoltageViolation']]])\n",
    "obs['FlowViolation'] = th.as_tensor([[obs['FlowViolation']]])\n",
    "\n",
    "\n",
    "action, values, log_probs = model.policy.forward(obs)\n",
    "#print(obs['loss'])\n",
    "DCKTOBJ=take_action(DSSCKTOBJ,action)\n",
    "OBS=get_state(DSSCKTOBJ,G_INIT)\n",
    "print(\"The loss is:\",OBS['loss'])\n",
    "print(\"The topology violation status is:\",OBS['TopologicalConstr'])\n",
    "print(\"The voltage violation status is: \",OBS['VoltageViolation'])\n",
    "print(\"The branch flow violation status is: \",OBS['FlowViolation'])\n",
    "print(\"The amount of unserved energy is:\",OBS['Unserved Energy']*1000*25)\n",
    "#print(\"The voltage violation status is:\",OBS['VoltageViolation'])\n",
    "print(\"The convergence status is:\",OBS['Convergence'])\n",
    "print(\"For Unity Load: The Optimal Configuration is :.{}\",action)\n",
    "end = time.time()\n",
    "print(\"Run time [s]: \",end-start)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DL_project",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
